import re
import json

import jieba
import yaml


class preprocessText:
    def __init__(self, config):
        jieba.load_userdict(config['argot_path'])
        with open(config['translation_path']) as fp:
            self.translation = json.load(fp)

    def pre_caption(self, caption, max_words=30):
        caption = re.sub(
            r"([,.'!?\"()*#:;~])",
            '',
            caption.lower(),
        ).replace('-', ' ').replace('/', ' ').replace('<person>', 'person')

        caption = re.sub(
            r"\s{2,}",
            ' ',
            caption,
        )
        caption = caption.rstrip('\n')
        caption = caption.strip(' ')

        # truncate caption
        caption_words = caption.split(' ')
        if len(caption_words) > max_words:
            caption = ' '.join(caption_words[:max_words])

        return caption

    def replace_argot(self, text):
        words = jieba.cut(text)
        new_text = ''
        for i, word in enumerate(words):
            if word in self.translation.keys():
                new_text += self.translation[word]
            else:
                new_text += word

        return new_text


if __name__ == '__main__':
    config_file_path = r'D:\Project\multimoding\project\CLIP\classify\config.yaml'
    config = yaml.load(open(config_file_path, 'r', encoding='utf-8'), Loader=yaml.Loader)
    p = preprocessText(config)
    text = p.replace_argot("威力不怎么样 别外那位大神告知下这狗可以改p吗？")
    print(text)
